package edu.poly.cs9223.ass2.ec2.tools;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.Charset;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import ch.ethz.ssh2.Connection;
import ch.ethz.ssh2.ConnectionInfo;
import ch.ethz.ssh2.KnownHosts;
import ch.ethz.ssh2.ServerHostKeyVerifier;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;

public class JSSHTool {

	private static Integer SSH_CONNECTION_TIME_OUT = 10 * 60000;// (10min)
	
	Logger logger = LoggerFactory.getLogger(JSSHTool.class);
	JSch jsch = new JSch();
	Session session = null;

	private String getHostFingerPrint(String hostName) {

		String sshServerAddr = hostName;
		Connection sshConnection = new Connection(sshServerAddr);
		try {
			sshConnection.connect(new ServerHostKeyVerifier() {
				public boolean verifyServerHostKey(String hostname, int port,
						String serverHostKeyAlgorithm, byte[] serverHostKey)
						throws Exception {
					return true;
				}
			}, SSH_CONNECTION_TIME_OUT, SSH_CONNECTION_TIME_OUT);
			ConnectionInfo info = sshConnection.getConnectionInfo();
			String keyAlgorithm = info.serverHostKeyAlgorithm;
			byte[] hostKeys = sshConnection.getConnectionInfo().serverHostKey;
			String fingerprint = KnownHosts.createHexFingerprint(keyAlgorithm,
					hostKeys);
			logger.debug("The host fingerprint is {} ", fingerprint);
			return fingerprint;
		} catch (IOException e) {
		} finally {
			sshConnection.close();
		}
		return null;
	}
	public void sshLogin(String userName, String host, String privateKey,
			int port) throws Exception {

		logger.info("get {} finger print.....",host);
		
		//we have to wait until the hostFingerPrint info is ready,otherwise there will be errors while connecting
		String hostFingerPrint = null;
		while (hostFingerPrint == null)
			hostFingerPrint = this.getHostFingerPrint(host);
		
		jsch.addIdentity(privateKey);
		logger.debug("private key identity added ");
		session = jsch.getSession(userName, host, port);
		java.util.Properties config = new java.util.Properties();
		config.put("StrictHostKeyChecking", "no");
		session.setConfig(config);
		logger.info("connecting.....");
		session.connect();
		logger.info("session connected.....");
		logger.info("successfully connected to {} via SSH", host);
	}

	public void executeShellCommand(String cmd) {
		logger.debug("executing command {}", cmd);
		try {
			Channel channel = session.openChannel("exec");
			((ChannelExec) channel).setCommand(cmd);
			channel.setInputStream(System.in);
			channel.setOutputStream(System.out);
			((ChannelExec) channel).setErrStream(System.err);
			((ChannelExec) channel).setPty(true);
			InputStream in = channel.getInputStream();
			channel.connect();
			byte[] tmp = new byte[1024];
			while (true) {
				while (in.available() > 0) {
					int i = in.read(tmp, 0, 1024);
					if (i < 0)
						break;
					System.out.print(new String(tmp, 0, i));
				}
				if (channel.isClosed()) {
					System.out.println("exit-status: "
							+ channel.getExitStatus());
					break;
				}
				try {
					Thread.sleep(1000);
				} catch (Exception ee) {
				}
			}
			channel.disconnect();
		} catch (JSchException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	public void openShell()
	{
		Channel jschChannel = null;
		try {
			logger.debug("opening shell");
			jschChannel = session.openChannel("shell");
			jschChannel.setOutputStream(System.out);
			jschChannel.setInputStream(System.in);
			jschChannel.connect();
			while(true)
			{
				if(jschChannel.isClosed())
					break;
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	/**
	 * Note: if any command in the shell file requires interaction between user and the shell,please use executeShellCommand() instead
	 * @param shellScript
	 */
	public void executeShellScript(File shellScript) {
		InputStream    fis;
		BufferedReader br;
		String         line;

		try {
			fis = new FileInputStream(shellScript);
			br = new BufferedReader(new InputStreamReader(fis, Charset.forName("UTF-8")));
			while ((line = br.readLine()) != null) {
				this.executeShellCommand(line);
			}
		} catch (FileNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	public void uploadFile(String ftpRemoteDirectory, String fileToTransmit) {
		try {
			// Open the SFTP channel
			Channel channel = session.openChannel("sftp");
			channel.connect();
			ChannelSftp channelSFTP = (ChannelSftp) channel;
			channelSFTP.setPty(true);
			// Change to the remote directory
			logger.debug("Changing to FTP remote dir: {}", ftpRemoteDirectory);
			channelSFTP.cd(ftpRemoteDirectory);
			// Send the file to remote host
			File f = new File(fileToTransmit);
			logger.debug("Storing file as remote filename:{} ", f.getName());
			channelSFTP.put(new FileInputStream(f), f.getName());
		} catch (Exception e) {
			logger.error("Storing remote file failed. ");
			e.printStackTrace();
		}
	}

	public void closeSession() {
		session.disconnect();
	}

}
